use alloc::vec::Vec;

use crate::{
    ir::{graph::Graph, node::NodeId},
    pass::pass::{Pass, PassResult},
};

pub struct DeadCodeElimination {}

impl DeadCodeElimination {
    #[inline]
    fn kill_node(graph: &mut Graph, node_id: NodeId) {
        if node_id as usize >= graph.nodes.len() {
            return;
        }
        let node = &graph.nodes[node_id as usize];

        for &iv in &node.inputs {
            if let Some(cs) = graph.consumers.get_mut(iv as usize) {
                if let Some(pos) = cs.iter().position(|&x| x == node_id) {
                    cs.swap_remove(pos);
                }
            }
        }
        for &ov in &node.outputs {
            if let Some(p) = graph.producers.get_mut(ov as usize) {
                if p == &Some(node_id) {
                    *p = None;
                }
            }
        }
        let n = &mut graph.nodes[node_id as usize];
        n.inputs.clear();
        n.outputs.clear();
        n.attrs.clear();
        n.domain.clear();
        n.op_type.clear();
    }
}

impl Pass for DeadCodeElimination {
    fn name(&self) -> &'static str {
        "DeadCodeElimination"
    }

    fn run(&mut self, graph: &mut crate::ir::graph::Graph) -> super::pass::PassResult {
        let n = graph.nodes.len();
        if n == 0 {
            return PassResult::unchanged();
        }

        let mut live = alloc::vec![false; n];
        let mut stack: Vec<NodeId> = Vec::new();

        for &ov in &graph.outputs {
            if let Some(p) = graph.producers.get(ov as usize).and_then(|&x| x) {
                if !live[p as usize] {
                    live[p as usize] = true;
                    stack.push(p);
                }
            }
        }
        while let Some(node_id) = stack.pop() {
            let node = &graph.nodes[node_id as usize];
            for &iv in &node.inputs {
                if let Some(p) = graph.producers.get(iv as usize).and_then(|&x| x) {
                    if !live[p as usize] {
                        live[p as usize] = true;
                        stack.push(p);
                    }
                }
            }
        }

        let mut changed = false;
        for i in 0..n {
            let dead = !live[i] && !graph.nodes[i].op_type.is_empty();
            if dead {
                Self::kill_node(graph, i as NodeId);
                changed = true;
            }
        }

        PassResult { changed }
    }
}
